from motor.motor_asyncio import AsyncIOMotorDatabase
from models.uploadRecords import UploadRecord
from .base import BaseRepository

class UploadRecordRepository(BaseRepository[UploadRecord]):
    def __init__(self, db: AsyncIOMotorDatabase):
        super().__init__("upload_records", db)

    async def get_by_upload_id(self, upload_id: str) -> UploadRecord:
        return await self.collection.find_one({"upload_id": upload_id})   

    async def create_or_update(self, uploadRecord: UploadRecord) -> UploadRecord:
        """
        创建或更新数据
        """
        query = {"upload_id": uploadRecord.upload_id}
        uploadRecord_dict = uploadRecord.model_dump()

        # 使用 update_one 方法，upsert=True 表示如果不存在则插入，存在则更新
        await self.collection.update_one(
            query,
            {
                "$set": uploadRecord_dict,
            },
            upsert=True
        )

        # 返回更新后的文档
        return await self.get_by_upload_id(uploadRecord.upload_id)